import json
import random
from collections import defaultdict
import sys

data_p = sys.argv[1]


def get_source(item):
    """Identify the source of an item from possible keys."""
    if "datasource" in item:
        return item["datasource"]
    elif "source" in item:
        return item["source"]
    elif "metadata" in item and isinstance(item["metadata"], dict):
        return item["metadata"].get("dataset")
    return None  # No source identified


# Load the JSON data
with open(data_p, "r") as f:
    data = json.load(f)

# Group items by their source
groups = defaultdict(list)
for item in data:
    source = get_source(item)
    groups[source].append(item)

# Seed for reproducibility
random.seed(42)
filtered_data = []

# Process each group
for source, items in groups.items():
    # Shuffle to ensure random selection
    random.shuffle(items)

    # Calculate 50% of the group (at least 1 item if not empty)
    group_size = len(items)
    k = max(1, int(group_size * 0.3)) if group_size > 0 else 0

    # Select items and add to result
    filtered_data.extend(items[:k])

# Save the filtered data
sub_d = data_p.replace(".json", "_sub.json")
with open(sub_d, "w") as f:
    json.dump(filtered_data, f, indent=2, ensure_ascii=False)
